'''
Created on Sep 21, 2009

@author: mkiyer

This goal of this program is to produce a BED file containing a list of 
polyA sites.  The format of the BED file is as follows:

field
-----
chrom
chromStart - some position before the polyA site
chromEnd - some position after the polyA site
name - some arbitrarily chosen name to associate with this site.  this could
       be the gene name where the polyA site occurs
score - not used, but required by BED, so this will contain a '.'
strand - depending on the application, this parameter may be set differently,
         if associating polyA sites with genes, then this will be the strand
         that the gene occurs on
thickStart - the position of the polyA site
thickEnd - thickStart + 1

Remaining fields are not used at this time.

There are various ways to produce this type of BED file.  The user should
be able to choose from various polyA-databases and/or algorithms using the
command-line options or by importing and calling only the methods of 
interest.
'''

from bx.intervals.intersection import Interval, IntervalTree
from veggie.gene.gene import parse_refseq_genes
from veggie.genome.chrom import get_chrom_names
import os
import collections
import sys
import logging
from optparse import OptionParser

__path_to_tables = '/lab/mkiyer/refdb/ucsc_tables'

def _read_polyadb_table(fhd):
    '''
    the polyaDb table is downloaded directly from UCSC
    
    in this table the polya site is given in the chromStart, chromEnd fields
    '''
    for thisline in fhd:
        thisline = thisline.strip()
        if thisline == None:
            continue
        if thisline.startswith('#'):
            continue
        thisfields = thisline.split('\t')
        chrom = thisfields[1]
        start = int(thisfields[2])
        yield (chrom, start)

def _read_polyapredict_table(fhd):
    '''
    the format of this table (downloaded from UCSC) is:
    #bin    chrom   chromStart      chromEnd        name    score   strand  thickStart      thickEnd
    
    the polya site is in the thickStart field
    '''
    for thisline in fhd:
        thisline = thisline.strip()
        if thisline == None:
            continue
        if thisline.startswith('#'):
            continue
        thisfields = thisline.split('\t')
        chrom = thisfields[1]
        if chrom not in get_chrom_names():
            continue
        start = int(thisfields[7])
        yield (chrom, start)

def gen_polya_sites(use_polyadb, polyadb_fhd, use_polyapredict, polyapredict_fhd):
    '''
    returns a the union of all polya sites in polyaDb and polyPredict
    '''
    polya_sites = set([])
    if use_polyadb:
        # read polya sites from polyadb
        polyadb_sites = list(_read_polyadb_table(polyadb_fhd))
        logging.debug('polyadb: read %d sites' % len(polyadb_sites))
        polya_sites.update(polyadb_sites)
    if use_polyapredict:
        # read polya sites from polyapredict
        polyapredict_sites = list(_read_polyapredict_table(polyapredict_fhd))
        logging.debug('polyaPredict: read %d sites' % len(polyapredict_sites))
        polya_sites.update(polyapredict_sites)
    logging.debug('read total of %d polya sites' % len(polya_sites))
    return polya_sites

def build_interval_trees(genes):
    # build interval tree of 3' UTR exons
    utr_trees = collections.defaultdict(lambda: IntervalTree())
    gene_trees = collections.defaultdict(lambda: IntervalTree())
    
    # keep track of all UTRs added so the same UTR is not added more than once 
    utrs = set([])
    for g in genes:
        # add a polya_sites list attribute to each gene
        g.utr3_polya_sites = []
        # add gene to gene tree
        gene_trees[g.chrom].insert_interval(Interval(g.tx_start, g.tx_end, value=g))
        # add 3'UTR to utr tree
        for e in g.utr3_exons:
            start, end = e[0], e[1]
            if (g.chrom, e[0], e[1]) not in utrs:
                interval = Interval(start, end, value=g)
                utr_trees[g.chrom].insert_interval(interval)
                utrs.add((g.chrom, e[0], e[1]))
    return gene_trees, utr_trees

def filter_overlapping_genes(gene_trees, chrom, start, end):
    '''
    Unfortunately, there are many pairs of genes where the 3'UTR of one
    gene overlaps with the 5' UTR of another, and vice versa.  In these
    cases it is not easy to disambiguate reads mapping to the overlapping
    region.
        
    the gene interval tree is created with value=Gene, so we can ensure
    that all the genes that overlap the specified interval are isoforms
    of the same gene by checking the gene symbol
    
    return True if this UTR has only one gene, False otherwise
    '''
    # TODO: this algorithm is very strict.  sometimes genes are named differently
    # and overlap and are still functionally the same.  therefore this is too 
    # depending on proper gene symbols being assigned.
    #
    # Another approach would be to insist that each overlapping gene it itself
    # overlapping with all other hits and has the same strand orientation as all
    # other hits.  This might rescue some legitimately overlapping genes
    symbols = set([x.value.symbol for x in gene_trees[chrom].find(start, end)])
    if len(symbols) > 1:
        return False
    return True

def apply_interval_constraints(gene, start, end, polya, 
                               interval_min, interval_max,
                               utr_extend):    
    if not utr_extend:
        # the interval is too small to be considered testable and
        # extensions are not enabled
        if (polya - start < interval_min) or (end - polya < interval_min):
            return None, None
        
    if gene.strand == '+':    
        if (polya - start < interval_min):
            # TODO: cannot extend in 5' direction yet because need to incorporate
            # previous exons
            return None, None
        if (end - polya < interval_min):
            # extend the UTR so that coverage can be measured
            end = polya + interval_min
    else:
        if (polya - start < interval_min):
            start = polya - interval_min
        if (end - polya < interval_min):
            # TODO: cannot extend in 5' direction yet because need to incorporate
            # previous exons
            return None, None
        
    if interval_max != 0:
        if (polya - start) > interval_max:
            start = polya - interval_max
        if (end - polya) > interval_max:
            end = polya + interval_max
    return start, end

def match_genes_with_utrs(gene_trees,
                          utr_trees,
                          polya_sites, 
                          filter_overlapping=True,
                          use_utr=True,
                          use_window=False,
                          window_size=200,
                          interval_min=0,
                          interval_max=0,
                          utr_extend=False):
    '''
    Given a pointer to the refflat gene table and a list of polya_sites, 
    add a list of possible polya sites on the 3'UTR of each gene.  The 
    list will be stored as the attribute g.utr3_polya_sites.
    
    Each item of the list has the form (start, end, polya), where
    [start, end) is the interval of the entire UTR, and 'polya' is
    the position of the polyA site in the interval
    '''
    assert use_window != use_utr

    # stats
    polyas_used = 0
    failed_interval_min = 0
    failed_overlapping = 0
    failed_utr_splicing = 0

    # match polya sites with genes
    for chrom, polya in polya_sites:
        # intersect the polyA site with the interval tree to find all UTRs
        # that have this polyA site inside them
        hits = utr_trees[chrom].find(polya, polya+1)
        passed_filters = False
        
        for hit in hits:            
            # get gene
            g = hit.value
            
            start, end = apply_interval_constraints(g, hit.start, hit.end, polya, 
                                                    interval_min, interval_max,
                                                    utr_extend)
            if (start == None) or (end == None):
                failed_interval_min += 1
                continue

            # need to ensure that all UTRs with this polyA site are from the
            # same gene            
            if filter_overlapping and not \
                filter_overlapping_genes(gene_trees, chrom, start, end):
                # skip this polyA site because it resides within more than one gene UTR
                failed_overlapping += 1
                continue


            # TODO: currently cannot handle genes with splicing in the 3'UTR
#            if len(g.utr3_exons) > 1:
#                failed_utr_splicing += 1
#                continue
            passed_filters = True
            g.utr3_polya_sites.append((start, end, polya))

        if passed_filters == True:
            polyas_used += 1

    print '%d/%d polya sites were used' % (polyas_used, len(polya_sites))
    print '%d intervals < interval_min' % (failed_interval_min)
    print '%d intervals skipped due to overlapping with >1 gene' % (failed_overlapping)
    print "%d intervals skipped due to splicing in 3'UTR" % (failed_utr_splicing)


def write_utr_polya_sites(genes, outfhd):
    '''
    Writes UTR polya sites to file in format:
    
        string  geneName;           "Gene symbol"
        string  name;               "Accession"
        string  chrom;              "Chromosome name"
        char[1] strand;             "+ or - for strand"
        uint    utr_start;          "UTR start"
        uint    utr_end;            "UTR end"
        uint    utr_cdsStart;       "Coding exon region start"
        uint    utr_cdsEnd;         "Coding exon region end"
        uint    exonCount;          "Number of UTRs"
        uint[exonCount] exonStarts; "UTR start positions"
        uint[exonCount] exonEnds;   "UTR end positions"
        uint[exonCount] polyA;      "polya site positions"

    '''
    # write the list of polya sites for each gene to a file:
    
    # output header
    outfhd.write('\t'.join(['#symbol',
                            'accession',
                            'chrom',
                            'strand',
                            'utr_tx_start',
                            'utr_tx_end',
                            'utr_cds_start',
                            'utr_cds_end',
                            'polya_count',
                            'polya_starts',
                            'polya_ends',
                            'polya_sites']))
    outfhd.write('\n')
    
    for g in genes:
        if len(g.utr3_polya_sites) == 0:
            starts = [g.utr3[0]]
            ends = [g.utr3[1]]
            polya_sites = [(g.utr3[1] if g.strand == '+' else g.utr3[0])]
        else:
            starts = [x[0] for x in g.utr3_polya_sites]
            ends = [x[1] for x in g.utr3_polya_sites]
            polya_sites = [x[2] for x in g.utr3_polya_sites]

        thisline = '\t'.join([g.symbol,
                              g.acc,
                              g.chrom,
                              g.strand,
                              str(g.utr3[0]),
                              str(g.utr3[1]),
                              str(g.utr3_cds[0]),
                              str(g.utr3_cds[1]),
                              str(len(g.utr3_polya_sites)),
                              ','.join(map(str,starts)),
                              ','.join(map(str,ends)),
                              ','.join(map(str,polya_sites))])
        outfhd.write('%s\n' % thisline)


if __name__ == '__main__':
    optionparser = OptionParser("usage: %prog [--window | --utrs] [options]")
    optionparser.add_option('--path-to_tables', dest='path_to_tables',
                            default=__path_to_tables,
                            help="path (absolute or relative) to the ucsc"
                                 " tables polyaDb.txt, polyaPredict.txt, and"
                                 " refFlat.txt.  These tables must reside"
                                 " in the same directory [default=%default]")
    optionparser.add_option('--polyadb', action='store', default=True, 
                            dest="polyadb",
                            help="use the 'polyaDb' list of polya sites default=[%default]")
    optionparser.add_option('--polyapredict', action='store', default=True, 
                            dest="polyapredict",
                            help="use the 'polyaPredict' list of polya sites default=[%default]")
    optionparser.add_option('--nofilter-overlapping', action='store', default=True,
                            dest='overlap_filter',
                            help="do not remove polya sites that overlap more than one gene default=[%default]")    
    optionparser.add_option('-w', '--window', action='store_true', default=False, 
                            dest='window',
                            help="make table with window on either side of polya site default=[%default]")
    optionparser.add_option('--window-size', default=200, type="int", dest='window_size',
                            help="window size to use when in '-window' mode default=[%default]")
    optionparser.add_option('--utrs', action='store_true', default=False,
                            dest='utrs',
                            help="make table using 3'UTR exon of genes")
    optionparser.add_option('--utr-min-length', default=0, type="int", dest='interval_min',
                            help="minimum allowable distance from the polyA"
                                 " site to the transcript end.  polyA sites"
                                 " too close to the start/end of a transcript"
                                 " will not result in reliable coverage" 
                                 " measurements")
    optionparser.add_option('--utr-max-length', default=0, type="int", dest='interval_max',
                            help="maximum allowable distance from the polyA"
                                 " site to the transcript end.  large intervals"
                                 " may be subject to excessive coverage" 
                                 " variation and result in unreliable results")
    optionparser.add_option('--utr-extend', action="store_true", 
                            default=False, dest='utr_extend',
                            help="minimum allowable distance from the polyA"
                                 " site to the transcript end.  polyA sites"
                                 " too close to the start/end of a transcript"
                                 " will not result in reliable coverage" 
                                 " measurements")
    optionparser.add_option("-o", "--output", dest="outfile",
                            help="output file [default: %default]")    
    (options, args) = optionparser.parse_args()

    # build list of polya sites
    path_to_tables = options.path_to_tables
    polyadb_fhd = open(os.path.join(path_to_tables, 'polyaDb.txt'))
    polyapredict_fhd = open(os.path.join(path_to_tables, 'polyaPredict.txt'))
    polya_sites = gen_polya_sites(options.polyadb, polyadb_fhd,
                                  options.polyapredict, polyapredict_fhd)
    polyadb_fhd.close()
    polyapredict_fhd.close()

    # read list of genes
    refflat_fhd = open(os.path.join(path_to_tables, 'refFlat.txt'))
    genes = list(parse_refseq_genes())
    refflat_fhd.close()

    # get interval trees for all genes and for all 3' UTRs    
    gene_trees, utr_trees = build_interval_trees(genes)

    # try to associate each polya site with overlying genes using 
    # the constraints specified in the command-line parameters
    match_genes_with_utrs(gene_trees,
                          utr_trees,
                          polya_sites, 
                          filter_overlapping=True,
                          use_utr=options.utrs,
                          use_window=options.window,
                          window_size=options.window_size,
                          interval_min=options.interval_min,
                          interval_max=options.interval_max,
                          utr_extend=options.utr_extend)
        
    # write to file
    if options.outfile is None:
        outfhd = sys.stdout
    else:
        outfhd = open(options.outfile, 'w')    
    write_utr_polya_sites(genes, outfhd)
    outfhd.close()
